#!/usr/bin/env python3
import subprocess
from dataclasses import dataclass
import os
import copy


TASK = '203'
NNUNET_DATASET = '/home/yusongli/Templates/yunet'
ENV = {
    'nnUNet_n_proc_DA': '8',
    'nnUNet_raw': f'{NNUNET_DATASET}/nnUNet_raw',
    'nnUNet_preprocessed': f'{NNUNET_DATASET}/nnUNet_preprocessed',
    'nnUNet_results': f'{NNUNET_DATASET}/nnUNet_results',
    'CUDA_DEVICE_ORDER': 'PCI_BUS_ID',
}


@dataclass
class Group():
    cuda: str
    fold: str
    trainer: str


def train(groups):
    processes = []
    for group in groups:
        cmd = [
            '/home/yusongli/.Local/bin/anaconda3/bin/conda',
            'run', '-n', 'brats21', 'nnUNet_train', '3d_fullres', group.trainer, TASK, group.fold, '--npz'
        ]
        env = copy.deepcopy(ENV)
        env['CUDA_VISIBLE_DEVICES'] = group.cuda
        p = subprocess.Popen(cmd, env=env)
        processes.append(p)

    for p in processes:
        p.wait()


def predict(groups):
    processes = []
    for group in groups:
        _base_outputfolder = f'/home/yusongli/Templates/yunet/nnUNet_results/Dataset002_C_intensity1500_roi2.0/{group.trainer}__nnUNetPlans__3d_fullres/fold_{group.fold}'
        cmd = [
            '/home/yusongli/.Local/bin/anaconda3/bin/conda',
            'run', '-n', 'yunet', 'nnUNetv2_predict',
            '-i', f'/home/yusongli/Templates/yunet/nnUNet_raw/Dataset002_C_intensity1500_roi2.0/fold_{group.fold}/imagesTr',
            '-o', f'{_base_outputfolder}/validation',
            '-d', TASK,
            '-tr', group.trainer,
            '-c', '3d_fullres',
            '-f', group.fold,
            '-chk', f'{_base_outputfolder}/checkpoint_best.pth',
        ]
        env = copy.deepcopy(ENV)
        env['CUDA_VISIBLE_DEVICES'] = group.cuda
        p = subprocess.Popen(cmd, env=env)
        processes.append(p)

    for p in processes:
        p.wait()


if __name__ == '__main__':

    # WARNING: Don't use this, instead, use `./_.sh`.
    # WARNING: Don't use this, instead, use `./_.sh`.
    # WARNING: Don't use this, instead, use `./_.sh`.

    pass

    # groups = [
    #     Group(cuda='4', fold='0', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet'),
    #     Group(cuda='4', fold='2', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet'),
    #     Group(cuda='4', fold='4', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet'),
    #
    #     Group(cuda='5', fold='0', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1'),
    #     Group(cuda='5', fold='2', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1'),
    #     Group(cuda='5', fold='4', trainer='nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1'),
    # ]
    # train(groups)
    # # predict(groups)
